import { snippetCompletion } from '@codemirror/autocomplete';
import {
  LRLanguage,
  LanguageSupport,
  StreamLanguage,
  delimitedIndent,
  foldInside,
  foldNodeProp,
  indentNodeProp,
  syntaxHighlighting,
} from '@codemirror/language';
import { jinja2 } from '@codemirror/legacy-modes/mode/jinja2';
import { parseMixed } from '@lezer/common';
import { classHighlighter, styleTags, tags as t } from '@lezer/highlight';
import { vscodeDarkInit } from '@uiw/codemirror-theme-vscode';
import { parser } from './syntax.grammar.js';

export const BAMLLanguage = LRLanguage.define({
  parser: parser.configure({
    wrap: parseMixed((node) => {
      if (node.name !== 'PromptExprContents' && node.name !== 'JinjaExpr') {
        return null;
      }

      return {
        parser: StreamLanguage.define(jinja2).parser,
      };
    }),
    props: [
      indentNodeProp.add({
        Decl: delimitedIndent({ closing: '}', align: true }),
        // not sure the except is doing anything.
        PromptExpr: delimitedIndent({ closing: '#"', align: true }),
      }),
      foldNodeProp.add({
        ClassDecl: foldInside,
      }),
      styleTags({
        'Decl/*/"{" Decl/*/"}"': t.brace,
        'ValueExpr/"{" ValueExpr/"}"': t.brace,

        TestDecl: t.keyword,
        'TestDecl/IdentifierDecl': t.typeName,

        TypeBuilderDecl: t.keyword,
        TypeBuilderKeyword: t.keyword,

        EnumDecl: t.keyword,
        'EnumDecl/IdentifierDecl': t.typeName,
        'EnumDecl/EnumValueDecl/IdentifierDecl': t.propertyName,

        ClassDecl: t.keyword,
        'ClassDecl/IdentifierDecl': t.typeName,

        TypeAliasDecl: t.keyword,
        'TypeAliasDecl/IdentifierDecl': t.typeName,

        DynamicDecl: t.keyword,
        'DynamicDecl/IdentifierDecl': t.typeName,

        ClientDecl: t.keyword,
        'ClientDecl/IdentifierDecl': t.typeName,

        FunctionDecl: t.keyword,
        'FunctionDecl/IdentifierDecl': t.function(t.variableName),
        'FunctionArg/IdentifierDecl': t.variableName,

        'ClassField/IdentifierDecl': t.propertyName,
        
        // Constructor styling to match class declarations
        'ObjectConstructor/IdentifierDecl': t.typeName,
        'ObjectField/IdentifierDecl': t.propertyName,

        // Control flow keywords
        IfKeyword: t.controlKeyword,
        ElseKeyword: t.controlKeyword, 
        WhileKeyword: t.controlKeyword,
        
        // Operators
        BinaryOperator: t.operator,
        AssignmentOperator: t.operator,
        NumericLiteral: t.number,
        QuotedString: t.string,
        UnquotedString: t.string,
        'AttributeValue/UnquotedAttributeValue': t.string,
        FieldAttribute: t.attributeName,

        'FieldAttribute/@': t.attributeName,
        'FieldAttribute/IdentifierDecl': t.attributeName,
        BlockAttribute: t.attributeName,
        'BlockAttribute/IdentifierDecl': t.attributeName,

        'SimpleTypeExpr/IdentifierDecl': t.typeName,

        variable: t.controlKeyword,
        PromptExpr: t.string,
        'PromptExprNonJinja/...': t.string,
        'PromptExprNonJinja/PromptExprContents/...': t.operator,

        'TemplateStringDecl/IdentifierDecl': t.attributeName,
        'TemplateStringDecl/PromptExpr': t.string,

        'TupleValue/IdentifierDecl': t.operator,

        TrailingComment: t.comment,
        MultilineComment: t.comment,
      }),
    ],
  }),
  languageData: {
    commentTokens: { line: '//' },
    closeBrackets: {
      brackets: ['(', '[', '"', '#"', '{'],
      stringPrefixes: ['#"'],
      wordChars: ['#', '"'],
    },
    snippetCompletion: true,
  },
});

const exampleCompletion = BAMLLanguage.data.of({
  autocomplete: [
    snippetCompletion('@alias(#"${one}"#)', { label: '@alias' }),
    snippetCompletion('@assert({{ this }})', { label: '@assert' }),
    snippetCompletion('@check(custom_check_name, {{ this }})', {
      label: '@check',
    }),
    snippetCompletion('@description(#"${}"#)', { label: '@description' }),
    snippetCompletion('class ${ClassName} {\n  ${property string}\n}', {
      label: 'class',
    }),
    snippetCompletion('enum ${EnumName} {\n  ${ONE}\n  ${TWO}\n}', {
      label: 'enum',
    }),
    snippetCompletion(
      'function MyFunction(${arg: string}) -> ${YourOutputType} {\n  ${client YourClient}\n  prompt #"\n    ${add some instructions}\n  "#\n}',
      { label: 'function' },
    ),
    snippetCompletion(
      'prompt #"\n  {{ _.role("user") }}\n  INPUT:\n  ---\n  {{ your-variable }}\n  ---\n  Response:\n"#',
      { label: 'prompt #"' },
    ),
    snippetCompletion('_.role("${role}")', { label: '_.role"' }),
    snippetCompletion('#"${mystring}"#', { label: '#"' }),
    snippetCompletion(
      'client<llm> GPT4 {\n  provider baml-openai-chat\n  options {\n    model gpt4  \n}}',
      {
        label: 'client<llm> GPT4',
      },
    ),
    snippetCompletion(
      'template_string MyString(${arg1}: string) #"\n  A jinja string\n"#',
      {
        label: 'template_string',
      },
    ),
    snippetCompletion('type ${AliasName} = ${TypeExpr};', {
      label: 'type alias',
    }),
  ],
});
export const theme: ReturnType<typeof vscodeDarkInit> = vscodeDarkInit({
  styles: [
    {
      tag: [t.variableName],
      color: '#dcdcaa',
    },
    {
      tag: [t.brace],
      color: '#569cd6',
    },
    {
      tag: [t.variableName, t.propertyName],
      color: '#d4d4d4',
    },
    {
      tag: [t.attributeName],
      color: '#c586c0',
    },
  ],
});

export function BAML() {
  return new LanguageSupport(BAMLLanguage, [
    exampleCompletion,
    syntaxHighlighting(classHighlighter),
  ]);
}
